package com.dsy.agent;

import cn.hutool.core.collection.CollUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;

@Component
public class ControlAgent extends ReActAgent {
    private static final Logger logger = LoggerFactory.getLogger(BaseAgent.class);
    @Autowired
    private ChatClient chatClient;
    @Autowired
    private ToolCallingManager toolCallingManager;
    private ChatResponse chatResponse;

    @Override
    public boolean think() {
        //拼接提示词，从历史记忆中
        List<Message> messages = getMessages();
        messages.add(new UserMessage(getNextStepPrompt()));
        Prompt prompt = new Prompt(messages);
        //调用ai大模型api，从响应中可以得到是否需要调用工具
        chatResponse = chatClient.prompt(prompt).call().chatResponse();
        if (!chatResponse.hasToolCalls()) {
            AssistantMessage output = chatResponse.getResult().getOutput();
            messages.add(output);
            return false;
        }
        return true;
    }

    @Override
    public String act() throws InterruptedException {
        //从历史记忆中构建提示词
        List<Message> messages = getMessages();
        Prompt prompt = new Prompt(messages);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);
        List<Message> toolExecMesList = toolExecutionResult.conversationHistory();

        setMessages(toolExecMesList);
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        // 判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> response.name().equals("doTerminate"));
        if (terminateToolCalled) {
            // 任务结束，更改状态
            setAgentState(Agent.AGENT_STATE_FINISH);
        }

        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() + " 返回的结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        logger.info(results);
        return results;

    }
}
